import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import pickle
from skimage.io import imread, imshow
import matplotlib.pyplot as plt

# 假设你已经有了砂岩截面图和分区图的路径
sandstone_image_path = 'Sandstone_1.tif'
partition_image_path = 'Sandstone_1_segment.tif'

# 读取图像
sandstone_image = imread(sandstone_image_path)
partition_image = imread(partition_image_path)

# 输出图像形状
print(f"图像形状: {sandstone_image.shape}")

# 确保图像大小相同
assert sandstone_image.shape[:2] == partition_image.shape[:2], "图像大小不匹配"

# 将砂岩截面图转换为特征矩阵
# 假设砂岩截面图是灰度图
X = sandstone_image.reshape(-1, 1)

# 将彩色分区图转换为标签向量
# 假设彩色分区图的每个像素颜色代表一个类别
# 这里我们简化处理，将每个唯一的颜色映射到一个唯一的类别标签
unique_colors, labels = np.unique(partition_image.reshape(-1, partition_image.shape[-1]), axis=0, return_inverse=True)

print("完成从砂岩截面图1及其对应分区中获取X和y")

# 检查X和y的长度是否一致
assert X.shape[0] == labels.shape[0], f"特征矩阵和标签向量的长度不一致: {X.shape[0]} != {labels.shape[0]}"

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, labels, test_size=0.2, random_state=42)
print("完成train_test_split")

# 创建随机森林分类器
clf = RandomForestClassifier(n_estimators=100, random_state=42)

# 训练模型
clf.fit(X_train, y_train)
print("完成随机森林模型clf的训练")

# 预测整个图像
y_pred = clf.predict(X)

# 将预测结果转换回图像形状
predicted_segmentation = y_pred.reshape(sandstone_image.shape[:2])

# 计算准确率
accuracy = accuracy_score(y_test, clf.predict(X_test))
print(f"准确率:  {accuracy}")

# 保存模型到硬盘
with open('sandstone_classifier.pkl', 'wb') as file:
    pickle.dump(clf, file)

print("已保存随机森林模型clf到硬盘")

# 显示预测分区灰度图
plt.imshow(predicted_segmentation, cmap='gray', vmin=0, vmax=np.max(predicted_segmentation))
plt.title('Predicted Segmentation')
plt.axis('off')
plt.show()